/*
 * Copyright 2023-2025, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 */

#include <arch_cpu_defs.h>
#include <boot/kernel_args.h>
#include <KernelExport.h>
#include <kernel.h>
#include <vm/vm.h>
#include <vm/vm_priv.h>
#include <vm/VMAddressSpace.h>

#include "LoongArchVMTranslationMap.h"

#define TRACE_VM_TMAP
#ifdef TRACE_VM_TMAP
#define TRACE(x...) dprintf(x)
#else
#define TRACE(x...) ;
#endif

// 全局变量
ssize_t gVirtFromPhysOffset = 0;
phys_addr_t sPageTable = 0;
char sPhysicalPageMapperData[sizeof(LoongArchVMPhysicalPageMapper)];

// 硬件特性标志初始化
uint32 LoongArchVMTranslationMap::fHwFeature = 0;

// 查找页表项
static loongarch_pte *
LookupPte(phys_addr_t pageTable, addr_t virtAdr)
{
    loongarch_pte *pte = (loongarch_pte *)VirtFromPhys(pageTable);
    for (int level = 3; level > 0; level--)
    {
        pte += VirtAddrPte(virtAdr, level);
        if (!pte->IsValid())
        {
            return NULL;
        }
        // 获取下一级页表的物理地址
        pte = (loongarch_pte *)VirtFromPhys(pte->PhysicalAddress());
    }
    pte += VirtAddrPte(virtAdr, 0);
    return pte;
}

// 创建页表项
static loongarch_pte *
LookupPteAlloc(addr_t virtAdr, bool alloc, kernel_args *args,
               phys_addr_t (*get_free_page)(kernel_args *))
{
    loongarch_pte *pte = (loongarch_pte *)VirtFromPhys(sPageTable);
    for (int level = 3; level > 0; level--)
    {
        pte += VirtAddrPte(virtAdr, level);
        if (!pte->IsValid())
        {
            if (!alloc)
                return NULL;
            page_num_t ppn = get_free_page(args);
            if (ppn == 0)
                return NULL;
            memset((loongarch_pte *)VirtFromPhys(B_PAGE_SIZE * ppn), 0, B_PAGE_SIZE);

            // 创建新的页表项
            pte->val = (ppn << PAGE_SHIFT) | kPteValid | kPteRead | kPteWrite | kPteExec;
        }
        // 获取下一级页表的物理地址
        pte = (loongarch_pte *)VirtFromPhys(pte->PhysicalAddress());
    }
    pte += VirtAddrPte(virtAdr, 0);
    return pte;
}

// 映射页面
static void
Map(addr_t virtAdr, phys_addr_t physAdr, uint64 flags, kernel_args *args,
    phys_addr_t (*get_free_page)(kernel_args *))
{
    TRACE("Map(0x%" B_PRIxADDR ", 0x%" B_PRIxADDR ")\n", virtAdr, physAdr);
    loongarch_pte *pte = LookupPteAlloc(virtAdr, true, args, get_free_page);
    if (pte == NULL)
        panic("无法分配页表");

    // 设置页表项
    pte->val = (physAdr & kPtePhysAddrMask) | flags | kPteValid | kPteAccessed | kPteDirty;

    // 刷新TLB
    asm volatile(
        "invtlb 0, $zero, %0\n"
        :
        : "r"(virtAdr)
        : "memory");
}

// 打印页表项标志
static void
WritePteFlags(uint32 flags)
{
    bool first = true;
    dprintf("{");
    for (uint32 i = 0; i < 32; i++)
    {
        if ((1 << i) & flags)
        {
            if (first)
                first = false;
            else
                dprintf(", ");

            switch (i)
            {
            case 0:
                dprintf("valid");
                break;
            case 1:
                dprintf("read");
                break;
            case 2:
                dprintf("write");
                break;
            case 3:
                dprintf("exec");
                break;
            case 4:
                dprintf("user");
                break;
            case 5:
                dprintf("global");
                break;
            case 6:
                dprintf("accessed");
                break;
            case 7:
                dprintf("dirty");
                break;
            case 8:
                dprintf("nocache");
                break;
            default:
                dprintf("%" B_PRIu32, i);
            }
        }
    }
    dprintf("}");
}

// 页表转储类
class PageTableDumper
{
private:
    uint64 firstVirt;
    uint64 lastVirt;
    phys_addr_t firstPhys;
    uint64 flags;
    int count;

    void Flush()
    {
        if (count == 0)
            return;

        dprintf("0x%016" B_PRIx64 " - 0x%016" B_PRIx64 " -> 0x%016" B_PRIx64 " - 0x%016" B_PRIx64 " ",
                firstVirt, lastVirt, firstPhys, firstPhys + (lastVirt - firstVirt));
        WritePteFlags(flags);
        dprintf(" (%" B_PRId32 " pages)\n", count);
    }

public:
    PageTableDumper()
        : firstVirt(0), lastVirt(0), firstPhys(0), flags(0), count(0)
    {
    }

    ~PageTableDumper()
    {
        Flush();
    }

    void Add(uint64 virt, phys_addr_t phys, uint64 _flags)
    {
        if (count > 0 && (virt != lastVirt + B_PAGE_SIZE || phys != firstPhys + (lastVirt - firstVirt + B_PAGE_SIZE) || flags != _flags))
        {
            Flush();
            count = 0;
        }

        if (count == 0)
        {
            firstVirt = virt;
            firstPhys = phys;
            flags = _flags;
        }

        lastVirt = virt;
        count++;
    }
};

// 转储页表
static void
DumpPageTable(phys_addr_t pageTable, uint64 virtAdrStart, uint64 virtAdrEnd)
{
    PageTableDumper dumper;

    for (uint64 virt = virtAdrStart; virt < virtAdrEnd; virt += B_PAGE_SIZE)
    {
        loongarch_pte *pte = LookupPte(pageTable, virt);
        if (pte != NULL && pte->IsValid())
        {
            dumper.Add(virt, pte->PhysicalAddress(), pte->val & 0x3ff);
        }
    }
}

// 转储内核页表
void DumpKernelPageTable()
{
    DumpPageTable(sPageTable, KERNEL_BASE, KERNEL_BASE + 0x400000000);
}

// 转储用户页表
void DumpUserPageTable(VMAddressSpace *addressSpace)
{
    LoongArchVMTranslationMap *map = (LoongArchVMTranslationMap *)addressSpace->TranslationMap();
    DumpPageTable(map->fPageTable, 0, USER_TOP);
}

// #pragma mark -

// 初始化虚拟内存转换映射
status_t
arch_vm_translation_map_init(kernel_args *args, VMPhysicalPageMapper **_physicalPageMapper)
{
    TRACE("arch_vm_translation_map_init\n");

    // 获取当前页表基址
    uint64 pageTable;
    asm volatile("csrrd %0, 0x19\n" : "=r"(pageTable));
    sPageTable = pageTable & ~0x3ffULL;

    // 初始化物理页映射器
    new (sPhysicalPageMapperData) LoongArchVMPhysicalPageMapper();
    *_physicalPageMapper = (VMPhysicalPageMapper *)sPhysicalPageMapperData;

    // 设置虚拟地址到物理地址的偏移
    gVirtFromPhysOffset = KERNEL_PMAP_BASE;

    // 检测硬件特性
    uint32 cpucfg2 = 0;
    asm volatile(
        "cpucfg %0, 2\n"
        : "=r"(cpucfg2));

    // 检查硬件是否支持访问位和脏位
    if (cpucfg2 & (1 << 12))
    {                                                 // 假设第12位表示硬件访问位支持
        LoongArchVMTranslationMap::fHwFeature |= 0x1; // 硬件访问位
    }
    if (cpucfg2 & (1 << 13))
    {                                                 // 假设第13位表示硬件脏位支持
        LoongArchVMTranslationMap::fHwFeature |= 0x2; // 硬件脏位
    }

    TRACE("arch_vm_translation_map_init: 页表基址 = 0x%" B_PRIxADDR ", 硬件特性 = 0x%x\n",
          sPageTable, LoongArchVMTranslationMap::fHwFeature);

    return B_OK;
}

// 创建虚拟内存转换映射
status_t
arch_vm_translation_map_create_map(bool kernel, VMTranslationMap **_map)
{
    phys_addr_t pageTable = 0;
    if (kernel)
    {
        // 内核使用当前页表
        pageTable = sPageTable;
    }
    else
    {
        // 用户空间需要分配新页表
        vm_page_reservation reservation;
        vm_page_reserve_pages(&reservation, 1, VM_PRIORITY_SYSTEM);
        vm_page *page = vm_page_allocate_page(&reservation, PAGE_STATE_WIRED | VM_PAGE_ALLOC_CLEAR);
        DEBUG_PAGE_ACCESS_END(page);
        pageTable = page->physical_page_number << PAGE_SHIFT;

        // 初始化用户空间页表
        // 复制内核空间页表项
        loongarch_pte *srcPte = (loongarch_pte *)VirtFromPhys(sPageTable);
        loongarch_pte *destPte = (loongarch_pte *)VirtFromPhys(pageTable);

        // 复制内核空间的页表项（高地址空间）
        for (int i = 256; i < kPageTableEntriesPerPage; i++)
        {
            destPte[i] = srcPte[i];
        }
    }

    *_map = new (std::nothrow) LoongArchVMTranslationMap(kernel, pageTable);
    if (*_map == NULL)
        return B_NO_MEMORY;

    return B_OK;
}